# -*- coding: utf-8 -*-
import os,sys
import imp
import xml.etree.ElementTree as xml
from cStringIO import StringIO
import maya.OpenMayaUI as OpenMayaUI
import pymel.core as pmc
import maya.cmds as cmds
import platform
try:
    imp.find_module('PySide2')
    import shiboken2 as sbk
    from PySide2.QtWidgets import *
    from PySide2.QtGui import *
    from PySide2.QtCore import *
    import pyside2uic as uic
except :
    from PySide.QtGui import *
    from PySide.QtCore import *
    import shiboken as sbk
    import pysideuic as uic

def getMayaWindow():
    ptr = OpenMayaUI.MQtUtil.mainWindow()
    return sbk.wrapInstance(long(ptr), QWidget)

def loadUiType(uiFile):
    parsed = xml.parse(uiFile)
    widget_class = parsed.find('widget').get('class')
    form_class = parsed.find('class').text
    with open(uiFile, 'r') as f:
        o = StringIO()
        frame = {}
        uic.compileUi(f, o, indent=0)
        pyc = compile(o.getvalue(), '<string>', 'exec')
        exec pyc in frame
        # Fetch the base_class and form class based on their type in the xml from designer
        form_class = frame['Ui_%s' % form_class]
        base_class = eval('%s' % widget_class)
    return form_class, base_class

ui_file = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'export_ani_UI.ui')
FormClass, BaseClass = loadUiType(ui_file)


class myWindow(FormClass, BaseClass):

    def __init__(self, parent=None):

        super(myWindow, self).__init__(parent)
        self.setupUi(self)

        self.setObjectName('MyUi')
        self.setWindowTitle('Bake_SetKey')
        self.move(700, 300)

        self.lineEdit_frame.setText('800')
        self.ptb_Bake_CV.clicked.connect(self.bake_CV)
        self.ptn_SetKeyFrame.clicked.connect(self.setKey)
        self.ptn_animExport.clicked.connect(self.export_ani)
        self.ptn_OpenFolder.clicked.connect(self.openPath)


    def bake_CV(self):
        self.int_start = int(self.lineEdit_strart_frame.text())
        self.int_end = int(self.lineEdit_end_frame.text())
        CVList = []
        allList = []
        num = 0
        parentConstraintList = pmc.ls(type='parentConstraint')
        pointConstraintList = pmc.ls(type = 'pointConstraint')
        orientConstraintList = pmc.ls(type='orientConstraint')
        aimConstraintList = pmc.ls(type='aimConstraint')
        scaleConstraintList = pmc.ls(type='scaleConstraint')

        allList = parentConstraintList + pointConstraintList + orientConstraintList + aimConstraintList + scaleConstraintList
        for x in allList:
            nodeName = x.nodeName()
            myBoolean = pmc.referenceQuery(nodeName, isNodeReferenced = True)
            if myBoolean:
                pass

            else:
                num = num + 1
                print num
                getParent = x.getParent().name()

                CVList.append(getParent)
        # pmc.select(CVList)
        # selObjects = pmc.ls(sl=1)
        # for x in CVList:
        try:
            cmds.bakeResults(CVList, t=(self.int_start, self.int_end), simulation=True)
        except:
            QMessageBox.warning(self, 'Error', u'提示: 当前场景没有被约束物体', QMessageBox.Yes)

    def setKey(self):

        self.int_Frame = int(self.lineEdit_3.text())
        pmc.currentTime(self.int_Frame, edit=True)

        selObjectList01 = pmc.ls(type='nurbsCurve')
        selObjectList02 = pmc.ls(type='nurbsSurface')
        selObjectList = selObjectList01 + selObjectList02
        cvShape = pmc.select(selObjectList, r=1)
        mySelObject = pmc.ls(sl=1)
        pmc.select(mySelObject, r=1)
        FKList = ['FKIKArm_L','FKIKArm_R']
        IKList = ['FKIKSpine_M', 'FKIKLeg_L', 'FKIKLeg_R']


                    
        defaultDict = {'IKLeg_R.swivel': 0.0, 'FKNeck_M.Global': 0.0, 'ctrlEye_L.pupil': 0.0, 'FKShoulder_R.Global': 0.0,
                          'IKSpine3_M.stretchy': 10.0, 'Fingers_L.cup': 0.0, 'ctrlMouth_M.upperRoll': 0.0,
                          'FKIKArm_L.FKIKBlend': 0.0, 'Fingers_R.middleCurl': 0.0, 'IKLeg_R.volume': 10.0,
                          'ctrlEye_R.blinkDriven': 0.0, 'IKLeg_R.toe': 0.0, 'FKIKLeg_L.FKIKBlend': 10.0,
                          'ctrlEye_R.pupilAdd': 0.0, 'Fingers_R.ringCurl': 0.0, 'Blink_L.Pupil': 0.0,
                          'FKIKArm_R.FKIKBlend': 0.0, 'IKLeg_L.swivel': 0.0, 'Fingers_L.indexCurl': 0.0, 'DYN_Ct.ear': 0,
                          'ctrlEye_L.upperLid': 0.0, 'Fingers_L.thumbCurl': 0.0, 'ctrlEye_R.eyeDriven': 0.0,
                          'IKLeg_R.Lenght2': 1.0, 'DYN_Ct.dynVis': 0, 'AimEye_L.follow': 10.0, 'IKLeg_L.rollAngle': 25.0,
                          'ctrlEye_R.upperLidOut': 0.0, 'AimEye_R.aim': 0.0, 'IKLeg_L.stretchy': 0.0,
                          'ctrlEye_R.lowerLidOutAuto': 0.0, 'Fingers_R.pinkyCurl': 0.0, 'ctrlEye_L.pupilAdd': 0.0,
                          'Vis_Con.xgen': 0, 'DYN_Ct.fkVis': 1, 'ctrlEye_L.upperLidOutAuto': 0.0, 'IKLeg_R.roll': 0.0,
                          'IKSpine3_M.Global': 0.0, 'ctrlEye_L.lowerLid': 0.0, 'FKHead_M.Global': 0.0,
                          'Fingers_R.indexCurl': 0.0, 'AimEye_R.follow': 10.0, 'ctrlEye_R.upperLid': 0.0,
                          'Blink_R.PupilAdd': 0.0, 'IKLeg_L.Lenght1': 1.0, 'IKLeg_L.Lenght2': 1.0, 'ctrlEye_R.pupil': 0.0,
                          'IKLeg_L.antiPop': 0.0, 'Blink_R.Pupil': 0.0, 'Fingers_R.thumbCurl': 0.0,
                          'ctrlEye_L.blinkDriven': 0.0, 'ctrlEye_L.eyeDriven': 0.0, 'Blink_L.PupilAdd': 0.0,
                          'IKLeg_R.stretchy': 0.0, 'PoleLeg_L.lock': 0.0, 'Fingers_L.pinkyCurl': 0.0, 'IKLeg_R.rollAngle': 25.0,
                          'IKLeg_L.toe': 0.0, 'RootX_M.CenterBtwFeet': 0.0, 'ctrlEye_L.upperLidOut': 0.0, 'Fingers_R.cup': 0.0,
                          'ctrlEye_L.lowerLidOut': 0.0, 'Fingers_L.middleCurl': 0.0, 'ctrlEye_L.lowerLidFollow': 3.0,
                          'RootX_M.legLock': 0.0, 'FKIKSpine_M.FKIKBlend': 10.0, 'ctrlEye_R.lowerLidFollow': 3.0,
                          'IKSpine3_M.volume': 10.0, 'IKLeg_R.antiPop': 0.0, 'ctrlMouth_M.stickyLips': 0.0,
                          'FKShoulder_L.Global': 0.0, 'IKSpine1_M.stiff': 5.0, 'PoleLeg_R.follow': 10.0,
                          'Fingers_R.spread': 0.0, 'ctrlMouth_M.lowerRoll': 0.0, 'IKSpine3_M.stiff': 5.0,
                          'Blink_L.blinkCenter': 0.2, 'IKLeg_L.volume': 10.0, 'IKLeg_R.Lenght1': 1.0, 'PoleLeg_R.lock': 0.0,
                          'ctrlEye_R.lowerLidOut': 0.0, 'Blink_R.blinkCenter': 0.2, 'ctrlEye_L.lowerLidOutAuto': 0.0,
                          'ctrlEye_R.upperLidFollow': 10.0, 'FKIKLeg_R.FKIKBlend': 10.0, 'AimEye_M.follow': 10.0,
                          'Fingers_L.spread': 0.0, 'ctrlEye_R.upperLidOutAuto': 0.0, 'AimEye_M.aim': 10.0, 'DYN_Ct.onOff': 0,
                          'AimEye_L.aim': 10.0, 'DYN_Ct.startFrame': 0, 'PoleLeg_L.follow': 10.0, 'Fingers_L.ringCurl': 0.0,
                          'IKLeg_L.roll': 0.0, 'ctrlEye_R.lowerLid': 0.0, 'ctrlEye_L.upperLidFollow': 10.0,
                          'ctrlMouth_M.teethWidth': 0.0}

        invalidDict = {'translateX':0, 'translateY':0, 'translateZ':0, 
                        'rotateX':0, 'rotateY':0, 'rotateZ':0, 
                        'scaleX':1, 'scaleY':1, 'scaleZ':1}

        for x in mySelObject:

            selTransform = x.getParent()

            if len(self.myJudge01(selTransform)) == 0:
                attrList = pmc.listAttr(selTransform, k=1)

                for at in attrList:
                    attrType = type(pmc.getAttr('%s.%s' % (selTransform.name(),at)))

                    if attrType == bool:
                        pass

                    else:

                        if at in invalidDict.keys():
                            if not pmc.getAttr('%s.%s' % (selTransform,at), l=1):
                                pmc.setAttr('%s.%s' % (selTransform,at), invalidDict[at])
                                pmc.setKeyframe('%s.%s'%(selTransform,at))


                        if '%s.%s' % (selTransform.name(),at) in defaultDict.keys():
                            if not pmc.getAttr('%s.%s' % (selTransform, at), l=1):
                                pmc.setAttr('%s.%s' % (selTransform,at), defaultDict['%s.%s' % (selTransform,at)])
                                pmc.setKeyframe('%s.%s'%(selTransform,at))

            else:
                pass

    def myJudge01(self, arg):
        self.myJudge = []

        objShape = arg.getChildren()[0]

        getObjEnabled = objShape.overrideEnabled.get()
        getObjType = objShape.overrideDisplayType.get()
        getObjVis  = pmc.getAttr('%s.visibility'%arg)
        if getObjEnabled == 1 and getObjType == 2 or getObjVis == False:
            self.myJudge.append(arg)

        allParentList = arg.getAllParents()
        
        for i in allParentList:
            getEnabled = i.overrideEnabled.get()
            getType = i.overrideDisplayType.get()
            getVis  = pmc.getAttr('%s.visibility'% i)
            if getEnabled == 1 and getType == 2 or getVis == False:
                
                self.myJudge.append(i)

            else:
               pass
        return self.myJudge

    def openPath(self):
        self.__fileBrower = QFileDialog()
        self.__fileBrower.resize(480, 270)
        self.directory_path = str( self.__fileBrower.getExistingDirectory())
        self.lineEdit_path.setText('%s' % (self.directory_path ))

    def export_ani(self):

        self.myPath = self.lineEdit_path.text()
        self.myFileName = self.lineEdit_file_name.text()

        mySceneName = cmds.file(q=1,sceneName=1)
        self.sceneDir, self.fileName = os.path.split(mySceneName)
        f_name, f_ex = os.path.splitext(self.fileName)
        print self.sceneDir
        if len(self.myPath) == 0:
            self.myPath = self.sceneDir
        if len(self.myFileName) == 0:
            self.myFileName = f_name
        sleCV = pmc.ls(sl=1)
        cmds.file(self.myPath + '/' + self.myFileName + '.anim', f=1, options="precision=17;hierarchy=below", type='animExport', pr=1, es=1)





def main():
    cmds.loadPlugin('mtoa', qt=1)
    cmds.loadPlugin('animImportExport', qt=1)
    if cmds.window('MyUi', exists=1):
        cmds.deleteUI('MyUi')
    e = myWindow(getMayaWindow())

    e.show()


if __name__ == '__main__':
    app = QApplication(sys.argv)

    sys.exit(app.exec_())